import re
from metagpt.actions import Action
from metagpt.logs import logger


class NaiveAddMainFunc(Action):
    #! Take as input a Rust code and return a Rust code with a main function added. Write the code to a file named `main.rs`.
    
    name: str = "NaiveAddMainFunc"
    
    ROLE_DESCRIPTION: str = """You are an expert in Rust programs. You are given a code snippet a new Rust learner translated from a reference C++ program. The code snippet does not have a main function. You should analyze the composition of the code and add a main function, which contains usage of all the defined functions ,classes and macro rules. The main function must be placed at the end of the code snippet. Response the whole code snippet with a main function added by yourself in this format and do not response with any other texts: 
    ```rust
        Your code here
    ```"""
    
    PROMPT_TEMPLATE: str = """{ROLE_DESCRIPTION}
    ## Reference C/C++ Program
    {cpp_code}
    ## Input Translated code
    {translated_code}
    ## Your Output
    """
    
    async def run(self, filenames: str):
        (cpp_file, translated_code_file) = filenames.split(";")
        with open(cpp_file, "r") as f:
            cpp_code = f.read()
        with open(translated_code_file, "r") as f:
            translated_code = f.read()

        pattern = r"fn main"
        match = re.search(pattern, translated_code, re.DOTALL)
        if not match:
            logger.info(f"No main function. Adding.")
            
            prompt = self.PROMPT_TEMPLATE.format(ROLE_DESCRIPTION=self.ROLE_DESCRIPTION, cpp_code=cpp_code, translated_code=translated_code)

            rsp = await self._aask(prompt)

            code_text = NaiveAddMainFunc.parse_code(rsp)
            # 写入 main.rs 文件
            with open(translated_code_file, "w") as f:
                f.write(code_text)
        else:
            logger.info(f"Main function detected")

        return filenames
    
    @staticmethod
    def parse_code(rsp):
        pattern = r"```rust(.*)```"
        match = re.search(pattern, rsp, re.DOTALL)
        code_text = match.group(1) if match else rsp
        return code_text
